# src/models.py

import os
from typing import Dict, Tuple
import torch
from torchvision import models

PRETRAINED_DIR = './pretrained'

# 各模型的分类头 key（用于过滤不匹配的头）
HEAD_KEYS: Dict[str, Tuple[str, str]] = {
    "efficientnet_b4": ("classifier.1.weight", "classifier.1.bias"),
    "resnet101": ("fc.weight", "fc.bias"),
}

def _get_tv_weights_enum(model_type):
    """返回 torchvision 的权重枚举；兼容不同版本的 torchvision。"""
    try:
        if model_type == "efficientnet_b4":
            from torchvision.models import EfficientNet_B4_Weights
            return EfficientNet_B4_Weights.IMAGENET1K_V1
        if model_type == "resnet101":
            from torchvision.models import ResNet101_Weights
            return ResNet101_Weights.IMAGENET1K_V1
    except Exception:
        # 旧版本兜底：传 None 走随机初始化（后续再用本地权重覆盖 backbone）
        return None
    return None

def get_weight_path(model_type):
    fname = {
        "efficientnet_b4": "efficientnet_b4_imagenet1k_v1.pth",
        "resnet101": "resnet101_imagenet1k_v1.pth",
    }.get(model_type)
    if not fname:
        raise ValueError(f"不支持的模型类型: {model_type}")
    os.makedirs(PRETRAINED_DIR, exist_ok=True)
    return os.path.join(PRETRAINED_DIR, fname)

def _replace_head(model_type, model, num_classes):
    """根据模型类型替换分类头。"""
    if model_type == "efficientnet_b4":
        in_features = model.classifier[1].in_features
        model.classifier[1] = torch.nn.Linear(in_features, num_classes)
    elif model_type == "resnet101":
        in_features = model.fc.in_features
        model.fc = torch.nn.Linear(in_features, num_classes)
    else:
        raise ValueError(f"不支持的模型类型: {model_type}")

def _safe_load_backbone(model, state_dict, model_type):
    """仅加载与当前结构匹配的参数；分类头尺寸不匹配则自动跳过。"""
    head_w, head_b = HEAD_KEYS[model_type]
    # 如果头尺寸不一致，直接从字典里剔除
    cur_sd = model.state_dict()
    to_pop = []
    for k in (head_w, head_b):
        if k in state_dict and k in cur_sd:
            if state_dict[k].shape != cur_sd[k].shape:
                to_pop.append(k)
    for k in to_pop:
        state_dict.pop(k, None)

    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if unexpected:
        print(f"⚠️ 未使用的权重键（unexpected）: {unexpected}")
    # missing 里通常包含被我们剔除的头，属正常情况
    return missing, unexpected

def load_torchvision_model(model_type, num_classes):
    # 1) 先用官方预训练（如可用）构建模型
    tv_weights = _get_tv_weights_enum(model_type)
    if model_type == "efficientnet_b4":
        model = models.efficientnet_b4(weights=tv_weights)
    elif model_type == "resnet101":
        model = models.resnet101(weights=tv_weights)
    else:
        raise ValueError(f"不支持的模型类型: {model_type}")

    # 2) 如果本地有缓存 pth，就作为“骨干覆盖”来加载（自动忽略不匹配的分类头）
    weight_path = get_weight_path(model_type)
    if os.path.exists(weight_path):
        print(f"🔎 发现本地权重：{weight_path}（将按“只加载骨干、忽略头”的方式加载）")
        try:
            ckpt = torch.load(weight_path, map_location="cpu")
            if isinstance(ckpt, dict) and "state_dict" in ckpt:
                ckpt = ckpt["state_dict"]
            _safe_load_backbone(model, ckpt, model_type)
        except Exception as e:
            print(f"⚠️ 本地权重载入失败，忽略并使用 torchvision 预训练：{e}")
    else:
        # 若没有本地文件且我们确实使用了官方权重，可将其缓存下来（注意：这是1000类头）
        if tv_weights is not None:
            try:
                torch.save(model.state_dict(), weight_path)
                print(f"✅ 已缓存官方预训练至 {weight_path}")
            except Exception as e:
                print(f"⚠️ 预训练缓存失败：{e}")

    # 3) 最后统一替换为下游任务的分类头
    _replace_head(model_type, model, num_classes)
    return model

def load_model(model_type, num_classes):
    if model_type in ['efficientnet_b4', 'resnet101']:
        return load_torchvision_model(model_type, num_classes)
    else:
        raise ValueError(f"不支持的模型类型: {model_type}")

# ====== 一键下载脚本（仅缓存官方预训练，不覆盖你的微调权重）======
if __name__ == '__main__':
    print("🚀 开始检测/缓存官方预训练 ...")
    for m in ["efficientnet_b4", "resnet101"]:
        try:
            print(f"\n>>> {m}")
            _ = load_torchvision_model(m, num_classes=1000)  # 构建 + 可能缓存
            print(f"✔️ {m} OK")
        except Exception as e:
            print(f"❌ {m} 失败: {e}")
    print("\n🎉 完成！")
